#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
未来传感器数据预测程序
调用训练好的模型，输入CSV历史数据，预测未来传感器数据
"""

import pandas as pd
import numpy as np
import os
import pickle
import warnings
from datetime import datetime, timedelta
import tensorflow as tf

warnings.filterwarnings('ignore')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class FutureSensorPredictor:
    """未来传感器数据预测器"""
    
    def __init__(self):
        """初始化预测器"""
        self.model = None
        self.scalers = None
        self.sensor_columns = [
            'S1_Temp', 'S2_Temp', 'S3_Temp', 'S4_Temp',
            'S1_Light', 'S2_Light', 'S3_Light', 'S4_Light',
            'S1_Sound', 'S2_Sound', 'S3_Sound', 'S4_Sound',
            'S5_CO2', 'S5_CO2_Slope', 'S6_PIR', 'S7_PIR'
        ]
        
    def load_model(self):
        """加载训练好的模型和缩放器"""
        model_path = '../model/saved_models/sensor_predictor.h5'
        scalers_path = '../model/saved_models/sensor_scalers.pkl'
        
        # 检查文件是否存在
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        if not os.path.exists(scalers_path):
            raise FileNotFoundError(f"缩放器文件不存在: {scalers_path}")
        
        try:
            # 加载模型
            print("加载训练好的模型...")
            self.model = tf.keras.models.load_model(model_path)
            print(f"模型加载成功: {model_path}")
            
            # 加载缩放器
            print("加载数据缩放器...")
            with open(scalers_path, 'rb') as f:
                self.scalers = pickle.load(f)
            print(f"缩放器加载成功: {scalers_path}")
            
        except Exception as e:
            print(f"加载失败: {e}")
            raise
    
    def prepare_time_features(self, datetime_series):
        """准备时间特征"""
        time_features = pd.DataFrame()
        time_features['Hour_sin'] = np.sin(2 * np.pi * datetime_series.dt.hour / 24)
        time_features['Hour_cos'] = np.cos(2 * np.pi * datetime_series.dt.hour / 24)
        time_features['DayOfWeek_sin'] = np.sin(2 * np.pi * datetime_series.dt.dayofweek / 7)
        time_features['DayOfWeek_cos'] = np.cos(2 * np.pi * datetime_series.dt.dayofweek / 7)
        time_features['IsWeekend'] = (datetime_series.dt.dayofweek >= 5).astype(int)
        
        return time_features.values
    
    def prepare_input_data(self, df, window_size=300):
        """
        准备输入数据
        
        参数:
        df: 输入数据DataFrame
        window_size: 时间窗口大小
        
        返回:
        标准化后的输入序列
        """
        print(f"🔧 处理输入数据，窗口大小: {window_size}")
        
        # 确保有DateTime列
        if 'DateTime' not in df.columns:
            df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%d-%m-%Y %H:%M:%S')
        
        # 排序并检查数据量
        df = df.sort_values('DateTime').reset_index(drop=True)
        
        if len(df) < window_size:
            raise ValueError(f"输入数据不足！需要至少{window_size}条记录，但只有{len(df)}条")
        
        # 取最后window_size条记录
        df = df.tail(window_size).reset_index(drop=True)
        
        print(f"使用数据时间范围: {df['DateTime'].min()} 到 {df['DateTime'].max()}")
        
        # 检查传感器列是否存在
        missing_cols = [col for col in self.sensor_columns if col not in df.columns]
        if missing_cols:
            raise ValueError(f"缺少传感器列: {missing_cols}")
        
        # 提取传感器数据
        sensor_data = df[self.sensor_columns].values.astype(np.float32)
        
        # 准备时间特征
        time_features = self.prepare_time_features(df['DateTime'])
        
        # 合并特征
        combined_features = np.hstack([sensor_data, time_features])
        
        # 标准化
        scaler = self.scalers['combined_scaler']
        normalized_data = scaler.transform(combined_features)
        
        # 重塑为模型输入格式 [1, window_size, features]
        input_sequence = normalized_data.reshape(1, window_size, -1)
        
        print(f"输入序列形状: {input_sequence.shape}")
        return input_sequence, df['DateTime'].max()
    
    def generate_future_times(self, last_datetime, num_predictions=96):
        """生成未来的时间序列"""
        future_times = []
        current_time = last_datetime + timedelta(minutes=30)  # 从下一个30分钟开始
        
        for i in range(num_predictions):
            future_times.append(current_time)
            current_time += timedelta(minutes=30)
        
        return future_times
    
    def apply_sensor_constraints(self, df):
        """应用传感器数据约束，确保预测值在合理范围内"""
        # 温度传感器：15-35°C
        temp_cols = [col for col in df.columns if 'Temp' in col]
        for col in temp_cols:
            df[col] = np.clip(df[col], 15, 35)
        
        # 光线传感器：0-1000
        light_cols = [col for col in df.columns if 'Light' in col]
        for col in light_cols:
            df[col] = np.clip(df[col], 0, 1000)
        
        # 声音传感器：0-1
        sound_cols = [col for col in df.columns if 'Sound' in col]
        for col in sound_cols:
            df[col] = np.clip(df[col], 0, 1)
        
        # CO2传感器：300-1000 ppm
        if 'S5_CO2' in df.columns:
            df['S5_CO2'] = np.clip(df['S5_CO2'], 300, 1000)
        
        # CO2斜率：-50到50
        if 'S5_CO2_Slope' in df.columns:
            df['S5_CO2_Slope'] = np.clip(df['S5_CO2_Slope'], -50, 50)
        
        # PIR传感器：0或1 (整数)
        pir_cols = [col for col in df.columns if 'PIR' in col]
        for col in pir_cols:
            df[col] = np.round(np.clip(df[col], 0, 1)).astype(int)
        
        return df
    
    def predict_future_sensors(self, csv_file, num_predictions=96, save_output=True):
        """
        预测未来传感器数据的主函数
        
        参数:
        csv_file: 输入CSV文件路径
        num_predictions: 预测的时间点数量 (默认96=48小时)
        save_output: 是否保存预测结果
        
        返回:
        预测结果DataFrame
        """
        print("=" * 60)
        print("未来传感器数据预测")
        print("=" * 60)
        
        # 检查输入文件
        if not os.path.exists(csv_file):
            raise FileNotFoundError(f"输入文件不存在: {csv_file}")
        
        print(f"输入文件: {csv_file}")
        print(f"预测时间点: {num_predictions}个 ({num_predictions*0.5:.1f}小时)")
        
        # 加载模型
        self.load_model()
        
        # 读取输入数据
        print(f"\n读取输入数据...")
        df = pd.read_csv(csv_file)
        print(f"读取完成，共{len(df)}条记录")
        
        # 准备输入数据
        input_sequence, last_datetime = self.prepare_input_data(df)
        
        # 预测
        print(f"\n开始预测未来{num_predictions}个时间点...")
        predictions = self.model.predict(input_sequence, verbose=0)
        print(f"预测完成，输出形状: {predictions.shape}")
        
        # 反标准化
        print("反标准化预测结果...")
        scaler = self.scalers['combined_scaler']
        denormalized_predictions = scaler.inverse_transform(predictions[0])
        
        # 分离传感器数据 (前16列是传感器，后5列是时间特征)
        sensor_predictions = denormalized_predictions[:, :len(self.sensor_columns)]
        
        # 生成未来时间
        future_times = self.generate_future_times(last_datetime, num_predictions)
        
        # 创建结果DataFrame
        print("构建预测结果...")
        results = []
        for i, future_time in enumerate(future_times):
            row = {
                'Date': future_time.strftime('%d-%m-%Y'),
                'Time': future_time.strftime('%H:%M:%S')
            }
            
            # 添加传感器预测值 (保留1位小数)
            for j, col in enumerate(self.sensor_columns):
                row[col] = round(float(sensor_predictions[i, j]), 1)
            
            results.append(row)
        
        results_df = pd.DataFrame(results)
        
        # 应用传感器约束
        results_df = self.apply_sensor_constraints(results_df)
        
        # 保存结果
        if save_output:
            # 确保保存目录存在
            save_dir = os.path.join('..', 'data', 'predict_sensordata')
            os.makedirs(save_dir, exist_ok=True)
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = os.path.join(save_dir, f"future_sensors_prediction_{timestamp}.csv")
            results_df.to_csv(output_file, index=False)
            print(f"预测结果已保存: {output_file}")
        
        # 显示预测结果统计
        self.print_prediction_summary(results_df, last_datetime)
        
        return results_df
    
    def print_prediction_summary(self, results_df, last_datetime):
        """打印预测结果统计"""
        print("\n" + "=" * 50)
        print("预测结果统计")
        print("=" * 50)
        
        print(f"历史数据结束时间: {last_datetime}")
        print(f"预测时间范围: {results_df['Date'].iloc[0]} {results_df['Time'].iloc[0]} 到 {results_df['Date'].iloc[-1]} {results_df['Time'].iloc[-1]}")
        print(f"预测时长: {len(results_df)*0.5:.1f}小时 ({len(results_df)}个时间点)")
        
        # 传感器统计
        print(f"\n 温度传感器预测范围:")
        temp_cols = [col for col in results_df.columns if 'Temp' in col]
        for col in temp_cols:
            print(f"   {col}: {results_df[col].min():.1f}°C - {results_df[col].max():.1f}°C (平均: {results_df[col].mean():.1f}°C)")
        
        print(f"\nCO2传感器预测:")
        if 'S5_CO2' in results_df.columns:
            print(f"   S5_CO2: {results_df['S5_CO2'].min():.0f} - {results_df['S5_CO2'].max():.0f} ppm (平均: {results_df['S5_CO2'].mean():.0f} ppm)")
        
        print(f"\nPIR传感器预测:")
        pir_cols = [col for col in results_df.columns if 'PIR' in col]
        for col in pir_cols:
            active_count = sum(results_df[col] > 0)
            print(f"   {col}: {active_count}/{len(results_df)} 次激活 ({active_count/len(results_df)*100:.1f}%)")
        
        print("=" * 50)

def main():
    """主函数"""
    print("未来传感器数据预测程序")
    
    # 获取输入文件
    print("\n请输入要预测的选项:")
    print("1. 使用默认数据文件 (../data/room_occupancy.csv)")
    print("2. 输入自定义文件路径")
    
    choice = input("请选择 (1 或 2，默认为1): ").strip()
    
    if choice == '2':
        csv_file = input("请输入CSV文件路径: ").strip()
    else:
        csv_file = '../data/room_occupancy.csv'
    
    # 获取预测时间点数量
    try:
        num_pred = input("请输入预测时间点数量 (默认96=48小时): ").strip()
        num_predictions = int(num_pred) if num_pred else 96
    except:
        num_predictions = 96
    
    # 创建预测器并运行
    try:
        predictor = FutureSensorPredictor()
        results = predictor.predict_future_sensors(
            csv_file=csv_file,
            num_predictions=num_predictions,
            save_output=True
        )
        
        print(f"\n预测完成！")
        print(f"成功预测了{len(results)}个时间点的传感器数据")
        
        # 显示前几行预测结果
        print(f"\n预测结果示例 (前5行):")
        display_cols = ['Date', 'Time', 'S1_Temp', 'S5_CO2']
        available_cols = [col for col in display_cols if col in results.columns]
        results_display = results[available_cols].head().copy()
        # 格式化数值显示（保留1位小数）
        for col in results_display.columns:
            if col not in ['Date', 'Time']:
                results_display[col] = results_display[col].round(1)
        print(results_display)
        
    except Exception as e:
        print(f"预测失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 